#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@File    : mngs_plot_pie.py
@Author  : Bing Liang
@Email   : believer19940901@gmail.com
@Date    : 2025/9/23 17:01
@Description : 宏基因组中绘制物种比例饼图
"""
from argparse import ArgumentParser, Namespace
from pathlib import Path
from typing import Any

import matplotlib as mpl

mpl.use("cairo")  # 必须在 pyplot 之前
import matplotlib.pyplot as plt
import pandas as pd


def load_data(in_file: Path) -> pd.DataFrame:
    """读取输入文件并返回 DataFrame"""
    try:
        df = pd.read_csv(in_file, sep="\t", encoding="utf-8")
    except UnicodeDecodeError:
        df = pd.read_csv(in_file, sep="\t", encoding="gbk")

    required_cols = {"S", "sReads"}
    if not required_cols.issubset(df.columns):
        raise ValueError(f"输入文件缺少必要列: {required_cols}")
    return df[["S", "sReads"]]


def prepare_data(df: pd.DataFrame, top: int) -> tuple[Any, Any]:
    """计算比例并聚合 top 物种"""
    total_reads = df["sReads"].sum()
    df["ratio"] = df["sReads"] / total_reads

    # 找出 top 物种阈值
    threshold = df["sReads"].nlargest(top).min()
    df["type"] = df.apply(lambda x: x["S"] if x["sReads"] >= threshold else "Other", axis=1)

    out_df = (
        df.groupby("type", as_index=False)
        .agg({"sReads": "sum", "ratio": "sum"})
        .sort_values(by="sReads", ascending=False)
    )

    # 标签格式：物种名 (reads, 百分比)
    out_df["label"] = out_df.apply(
        lambda x: f"{x['type'].rsplit('|', 1)[-1]} (n={x['sReads']}, {x['ratio']*100:.2f}%)",
        axis=1
    )
    return out_df, total_reads


def plot_pie(out_df: pd.DataFrame, total_reads: int, out_pdf: Path, out_png: Path):
    """绘制饼图并保存"""
    fig, ax = plt.subplots(figsize=(7, 6))
    wedges, _ = ax.pie(
        out_df["sReads"],
        startangle=140,
        textprops={'fontsize': 12}
    )

    ax.legend(
        wedges, out_df["label"],
        title=f"microbe (n={total_reads})",
        loc="center left",
        bbox_to_anchor=(1, 0, 0.5, 1)
    )
    ax.axis('equal')
    fig.tight_layout()

    # 保存
    out_pdf.parent.mkdir(parents=True, exist_ok=True)
    out_png.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(out_pdf)
    fig.savefig(out_png, dpi=300)
    plt.close(fig)


def main(args: Namespace):
    in_file = Path(args.in_file).absolute()
    out_pdf = Path(args.out_pdf).absolute()
    out_png = Path(args.out_png).absolute()

    if not in_file.is_file():
        raise FileNotFoundError(f"{in_file} is not a file")

    df = load_data(in_file)
    out_df, total_reads = prepare_data(df, args.top)
    plot_pie(out_df, total_reads, out_pdf, out_png)


if __name__ == "__main__":
    parser = ArgumentParser(description="绘制宏基因组物种比例饼图")
    parser.add_argument("-i", "--in_file", type=str, required=True, help="输入 TSV 文件，需包含 S 和 sReads 列")
    parser.add_argument("-t", "--top", type=int, default=1, help="显示前 N 个物种，其余归为 Other")
    parser.add_argument("-f", "--out_pdf", type=str, required=True, help="输出 PDF 文件路径")
    parser.add_argument("-g", "--out_png", type=str, required=True, help="输出 PNG 文件路径")
    main(parser.parse_args())
